from pycparser import c_ast, c_generator
from contextlib import contextmanager
from c_ast_decl import api

class ASTBuilder:
    class ScopeFrame:
        def __init__(self, stmts, parent_scope=None, current_loop=None, current_function=None):
            self.stmts = stmts          # 当前作用域的语句列表
            self.pending_ifs = []       # 未闭合的 if 节点栈
            self.current_loop = current_loop    # 当前循环上下文
            self.current_function = current_function  # 当前函数上下文
            self.symbols = {}           # 符号表：变量名 -> 类型
            self.parent_scope = parent_scope  # 父级作用域

        def lookup_variable(self, name):
            """ 在当前或父级作用域查找变量 """
            current = self
            while current:
                if name in current.symbols:
                    return current.symbols[name]
                current = current.parent_scope
            return None

    def __init__(self):
        self.root = c_ast.Compound(block_items=[])
        self.scope_stack = [self.ScopeFrame(self.root.block_items)]
        self.func_decls = dict()
        self.current_function = None
    
    def add_statement(self, stmt):
        """ 添加语句并自动注册变量到符号表 """
        current_scope = self.scope_stack[-1]
        current_scope.stmts.append(stmt)
        
        # 如果是变量声明，记录到符号表
        if isinstance(stmt, c_ast.Decl):
            var_name = stmt.name
            var_type = self._parse_decl_type(stmt)
            if var_name in current_scope.symbols:
                raise ValueError(f"Duplicate variable '{var_name}' in scope")
            current_scope.symbols[var_name] = var_type

    def _parse_decl_type(self, decl):
        """ 解析 Decl 节点的类型 """
        type_node = decl.type
        if isinstance(type_node, c_ast.TypeDecl):
            if isinstance(type_node.type, c_ast.IdentifierType):
                return ' '.join(type_node.type.names)
        return 'unknown'

    # ------------ 作用域管理 ------------
    @contextmanager
    def function_def(self, name, params, return_type="void"):
        # 创建函数定义节点
        # 构建参数列表
        param_decls = []
        for p_type, p_name in params:
            param_decl = c_ast.Decl(
                name=p_name,
                quals=[],
                storage=[],
                funcspec=[],
                align=0,
                type=c_ast.TypeDecl(
                    declname=p_name,
                    quals=[],
                    align=0,
                    type=c_ast.IdentifierType([p_type])
                ),
                init=None,
                bitsize=0
            )
            param_decls.append(param_decl)
        param_list = c_ast.ParamList(params=param_decls)

        # 构建函数声明
        func_decl = c_ast.FuncDecl(
            args=param_list,
            type=c_ast.TypeDecl(
                declname=name,
                quals=[],
                align=0,
                type=c_ast.IdentifierType([return_type])
            )
        )
        self.func_decls[func_decl.type.declname] = func_decl
        decl = c_ast.Decl(
            name=name,
            quals=[],
            storage=[],
            funcspec=[],
            align=0,
            type=func_decl,
            init=None,
            bitsize=0
        )

        # 创建函数定义节点
        func_def = c_ast.FuncDef(
            decl=decl,
            body=c_ast.Compound(block_items=[]),
            param_decls=None
        )
        self.scope_stack[-1].stmts.append(func_def)

        # 进入函数作用域
        func_scope = self.ScopeFrame(
            stmts=func_def.body.block_items,
            parent_scope=self.scope_stack[-1],
            current_function=func_def
        )
        self.current_function = func_def
        # 注册参数到符号表
        for p_type, p_name in params:
            func_scope.symbols[p_name] = p_type
        self.scope_stack.append(func_scope)
        try:
            yield func_def
        finally:
            self.scope_stack.pop()

    @contextmanager
    def for_loop(self, init, cond, step):
        # 创建 For 节点
        if isinstance(init, list):
            init = c_ast.DeclList(init)
        elif isinstance(init, c_ast.Decl):
            init = c_ast.DeclList([init])
        for_node = c_ast.For(init, cond, step, c_ast.Compound(block_items=[]))
        self.scope_stack[-1].stmts.append(for_node)
        
        # 进入循环作用域
        loop_scope = self.ScopeFrame(
            stmts=for_node.stmt.block_items,
            parent_scope=self.scope_stack[-1],
            current_loop=for_node
        )
        self.scope_stack.append(loop_scope)
        try:
            yield for_node
        finally:
            self.scope_stack.pop()
    @contextmanager
    def while_loop(self, cond):
        while_node = c_ast.While(
            cond=cond,
            stmt=c_ast.Compound(block_items=[])
        )
        self.scope_stack[-1].stmts.append(while_node)
        
        # 进入循环作用域
        loop_scope = self.ScopeFrame(
            stmts=while_node.stmt.block_items,
            parent_scope=self.scope_stack[-1],
            current_loop=while_node
        )
        self.scope_stack.append(loop_scope)
        try:
            yield while_node
        finally:
            self.scope_stack.pop()
    @contextmanager
    def if_then(self, cond):
        # 创建 If 节点
        if_node = c_ast.If(cond, c_ast.Compound(block_items=[]), None)
        self.scope_stack[-1].stmts.append(if_node)
        self.scope_stack[-1].pending_ifs.append(if_node)
        
        # 进入 then 分支作用域
        then_scope = self.ScopeFrame(
            stmts=if_node.iftrue.block_items,
            parent_scope=self.scope_stack[-1]
        )
        self.scope_stack.append(then_scope)
        try:
            yield if_node
        finally:
            self.scope_stack.pop()

    @contextmanager
    def if_else(self):
        # 处理 else 分支
        if_node = self.scope_stack[-1].pending_ifs.pop()
        if_node.iffalse = c_ast.Compound(block_items=[])
        
        # 进入 else 分支作用域
        else_scope = self.ScopeFrame(
            stmts=if_node.iffalse.block_items,
            parent_scope=self.scope_stack[-1]
        )
        self.scope_stack.append(else_scope)
        try:
            yield if_node
        finally:
            self.scope_stack.pop()

    def create_function_decl(self, name, return_type, params, **kwargs):
        function_decl = api.create_function_decl(name, return_type, params, **kwargs)
        self.func_decls[name] = function_decl.type
        # self.add_statement(function_decl)

        return function_decl
            
    def create_func_call(self, name, args):
        current_scope = self.scope_stack[-1]
        func_decl = self.func_decls[name]
        for idx in range(len(args)):
            arg = args[idx]
            if current_scope.lookup_variable(arg):
                args[idx] = c_ast.ID(arg)
            elif not isinstance(arg, c_ast.Node):
                arg_base_type = func_decl.args.params[idx].type
                if isinstance(arg_base_type, c_ast.PtrDecl):
                    arg_type = 'string'
                else:
                    arg_type = arg_base_type.type.names[0]
                args[idx] = c_ast.Constant(arg_type, arg)
        if name in self.func_decls:
            name = c_ast.ID(name)
        else:
            raise ValueError("函数未声明")
    
        func_call = c_ast.FuncCall(name=name, args=c_ast.ExprList(args))
        return func_call

    def add_func_call(self, name, args):
        func_call = self.create_func_call(name, args)
        self.add_statement(func_call)
        return func_call
    
    def add_ret(self, expr=None):
        current_scope = self.scope_stack[-1]
        if current_scope.lookup_variable(expr):
            expr = c_ast.ID(expr)
        else:
            ret_type = self.current_function.decl.type.type.type.names[0]
            expr = c_ast.Constant(type=ret_type, value=expr)
        ret = c_ast.Return(expr)
        self.add_statement(ret)
        return ret

    # ------------ 控制语句 ------------
    def add_break(self):
        """ 添加 break 语句（自动关联到最近的循环） """
        current = self.scope_stack[-1]
        while current:
            if current.current_loop:
                self.add_statement(c_ast.Break())
                return
            current = current.parent_scope
        raise RuntimeError("break must be inside a loop")
    
    def add_continue(self):
        """ 添加 break 语句（自动关联到最近的循环） """
        current = self.scope_stack[-1]
        while current:
            if current.current_loop:
                self.add_statement(c_ast.Continue())
                return
            current = current.parent_scope
        raise RuntimeError("break must be inside a loop")
    
    
    def create_binary_op(self, op, left, right):
        current_scope = self.scope_stack[-1]
        left_type = current_scope.lookup_variable(left)
        right_type = current_scope.lookup_variable(right)
        if not left_type and not right_type:
            raise ValueError("左右操作数不能都未定义")
        if left_type:
            left = c_ast.ID(left)
            if right_type:
                right = c_ast.ID(right)
            else:
                right = c_ast.Constant(left_type, right)
        else:
            left = c_ast.Constant(right_type, left)
            right = c_ast.ID(right)
        return c_ast.BinaryOp(op=op, left=left, right=right)

    def create_assignment(self, op, lvalue, rvalue):
        current_scope = self.scope_stack[-1]
        lvalue_type = current_scope.lookup_variable(lvalue)
        if not lvalue_type:
            raise ValueError("左操作未定义")
        lvalue = c_ast.ID(lvalue)
        if isinstance(rvalue, str):
            rvalue = c_ast.Constant(lvalue_type, rvalue)
        return c_ast.Assignment(op=op, lvalue=lvalue, rvalue=rvalue)
    
    def add_assignment(self, op, lvalue, rvalue):
        assignment = self.create_assignment(op, lvalue, rvalue)
        self.add_statement(assignment)
        return assignment

    def create_unary_op(self, op, expr):
        current_scope = self.scope_stack[-1]
        if not current_scope.lookup_variable(expr):
            raise ValueError("操作数未定义")
        expr = c_ast.ID(expr)        
        return c_ast.UnaryOp(op=op, expr=expr)
    
    def create_variable_decl(self, name, type_name, **kwargs):
        decl = api.create_variable_decl(name, type_name, **kwargs)
        self.add_statement(decl)
        return decl
    
    def generate_source(self):
        generator = c_generator.CGenerator()
        srcs = []
        for item in self.root.block_items:
            src = generator.visit(item)
            srcs.append(src)
        return ''.join(srcs)
